load("@bazel_skylib//rules:build_test.bzl", "build_test")
load(
    "//bazel:api.bzl",
    "mojo_binary",
    "mojo_test",
    "pkg_filegroup",
    "pkg_files",
    "strip_prefix",
)

_RUN_BENCHMARKS = [
    "algorithm/parallelize_overhead.mojo",
    "nn/bench_gather_reduce.mojo",
]

_EXTRA_CONSTRAINTS = {
    "algorithm/parallelize_overhead.mojo": select({
        "//:asan": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),
    "gpu/bench_allreduce.mojo": ["//:nvidia_gpu"],  # FIXME: KERN-1377
    "gpu/bench_ep_dispatch.mojo": ["//:nvidia_gpu"],  # Requires NVSHMEM shmem_module_init
    "nn/bench_gather_reduce.mojo": ["@platforms//:incompatible"],  # TODO(KERN-464): re-enable and remove this tag
    "gpu/bench_gemv.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2412
    "gpu/bench_pdl_copy.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2415
    "gpu/bench_elementwise.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_mha.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_bmm.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_kv_cache_ragged_rope.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_reduce.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_stencil.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_kv_cache_ragged_qkv_matmul.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_split_matmul_allreduce.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_matmul.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_kv_cache_ragged_flash_attention.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_kv_cache_ragged_flash_attention_paged.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "gpu/bench_grouped_matmul.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2477
    "gpu/bench_topk.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MSTDL-1880
}

_CPU_TESTS = glob(
    ["**/*.mojo"],
    exclude = [
        "**/*.yaml",
        "autotune/**",
        "demos/**",
        "gpu/**",
        "misc/**",
        "packages/**",
    ] + _RUN_BENCHMARKS,
)

[
    mojo_binary(
        name = src.split(".")[0],
        srcs = [src],
        enable_assertions = False,
        target_compatible_with = _EXTRA_CONSTRAINTS.get(src, []),
        deps = [
            "//max:comm",
            "//max:internal_utils",
            "//max:layout",
            "//max:nn",
            "@mojo//:stdlib",
        ],
    )
    for src in _CPU_TESTS
]

[
    build_test(
        name = src + ".test",
        targets = [src.split(".")[0]],
    )
    for src in _CPU_TESTS
]

[
    mojo_test(
        name = src + ".test",
        size = "enormous",
        srcs = [src],
        args = ["-t"],
        enable_assertions = False,
        target_compatible_with = _EXTRA_CONSTRAINTS.get(src, []),
        deps = [
            "//max:comm",
            "//max:internal_utils",
            "//max:layout",
            "//max:nn",
            "@mojo//:stdlib",
        ],
    )
    for src in _RUN_BENCHMARKS
]

[
    mojo_binary(
        name = src.split(".")[0],
        srcs = [src],
        data = select({
            "//:nvidia_gpu": ["@nvshmem_prebuilt//:host"],
            "//conditions:default": [],
        }),
        enable_assertions = False,
        env = select({
            "//:nvidia_gpu": {"MODULAR_SHMEM_LIB_DIR": "../+http_archive+nvshmem_prebuilt"},
            "//conditions:default": {},
        }),
        target_compatible_with = ["//:has_gpu"] + _EXTRA_CONSTRAINTS.get(src, []),
        deps = [
            "//max:comm",
            "//max:internal_utils",
            "//max:layout",
            "//max:nn",
            "//max:shmem",
            "@mojo//:stdlib",
        ],
    )
    for src in glob(["gpu/*.mojo"])
]

[
    build_test(
        name = src + ".test",
        tags = ["gpu"],
        targets = [src.split(".")[0]],
    )
    for src in glob(["gpu/*.mojo"])
]

pkg_files(
    name = "kernel_benchmark_input_files",
    srcs = glob(
        [
            "**/*.mojo",
            "**/*.yaml",
        ],
        exclude = [
            "demos/**",
            "misc/**",
            "packages/**",
        ],
    ),
    prefix = "kernel-benchmark",
    strip_prefix = strip_prefix.from_pkg(""),
    visibility = ["//visibility:public"],
)

pkg_filegroup(
    name = "kernel_benchmark_files",
    srcs = [
        ":kernel_benchmark_input_files",
        "//max/kernels/benchmarks/autotune:autotune_files",
    ],
    visibility = ["//visibility:public"],
)
